# 研究脉冲宽度与信噪比关系
from typing import Dict, List
import json
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from apps.cpm.rlm_context import RlmContext
from apps.cpm.cpm_engine import CpmEngine
from core.radar_core import RadarCore

class TauSnrAgent(object):
    def __init__(self):
        self.name = 'apps.cpm.agent.tau_snr_agent.TauSnrAgent'

    @staticmethod
    def run(context:RlmContext) -> Dict:
        
        '''
        作为整体任务第一步时，以用户自然语言为输入，输出为本步完成后的结果。
        '''
        print(f'tau ranges in radar equation')
        observes = TauSnrAgent._observe(context=context)
        plans = TauSnrAgent._reason(observes=observes, context=context)
        return TauSnrAgent._act(observes=observes, context=context, plans=plans)

    @staticmethod
    def _observe(context:RlmContext) -> Dict:
        observes = {}
        prompt = TauSnrAgent._gen_params_parse_prompt(user_query=context.current_input)
        resp = CpmEngine.infer(query=prompt)
        print(f'resp:\n{resp};')
        print(f'input:\n{resp[8:-3]};')
        params = json.loads(resp[8:-3])
        context.states['params'].update(params)
        observes.update(params)
        return observes
    
    @staticmethod
    def _reason(observes:Dict, context: RlmContext) -> List:
        plans = []
        return plans
    
    @staticmethod
    def _act(observes:Dict, context: RlmContext, plans: List) -> Dict:
        snr_db = np.linspace(int(observes['SNRs'][0]), int(observes['SNRs'][1]), int(observes['SNRs'][2]))
        # 将SNR转换为基数10
        snr = 10 ** (0.1 * snr_db)
        tau = RadarCore.radar_eq_tau(
            float(observes['p_t']), float(observes['f']), 
            float(observes['G']), float(observes['sigma']), 
            float(observes['T_e']), float(observes['F']), 
            float(observes['L']), [float(ri) for ri in observes['Rs']], snr
        )
        # 绘制图形
        plt.figure(1)
        plt.semilogy(snr_db, 1e6 * tau[0], 'k', label='R = 75 Km')
        plt.semilogy(snr_db, 1e6 * tau[1], 'k-.', label='R = 100 Km')
        plt.semilogy(snr_db, 1e6 * tau[2], 'k:', label='R = 150 Km')
        plt.grid(True)
        plt.legend()
        plt.xlabel('Minimum required SNR - dB')
        plt.ylabel('τ (pulsewidth) in μ sec')
        # 创建根目录
        base_folder = 'work/reports/e002'
        Path(base_folder).mkdir(parents=True, exist_ok=True)
        # 创建图片目录
        images_folder = f'{base_folder}/images'
        Path(images_folder).mkdir(parents=True, exist_ok=True)
        plt.savefig(f'{images_folder}/f00.png')
        # 创建md文件
        with open(f'{base_folder}/report.md', 'w', encoding='utf-8') as wfd:
            wfd.write(f'# 雷达设计：脉冲宽度与信噪比实验报告')
            wfd.write(context.current_input)
            wfd.write(f'\n![f00.png](./images/f00.png "")')
        # 规划执行步骤
        context.states['params'].update(observes)
        return {'result': 0, 'params': {}}
    
    @staticmethod
    def _gen_params_parse_prompt(user_query:str) -> str:
        template = '''
# 角色
假设你是资深雷达系统设计师

# 任务
你需要根据用户输入，提取出所有的参数，并转换为标准单位，以json格式输出。由于参数很多，而且参数表述和单位形式多样，你需要按照下面的步骤执行。

# 要求
你只需要给出json结果，不需要生成代码或进行解释。

# 步骤
1. 提取雷达工作频率f：名称可能为工作频率、载频，或者用f、f_c、fc、f0表示，单位可能为Hz（赫兹）、KHz（千赫兹）、MHz（兆赫兹）、GHz，如果是KHz需要将值乘以1000，MHz需要乘以e6，GHz需要乘以e9，提取值为v1；
2. 发射功率p_t：名称可能为发射功率，发射峰值功率，单位可能为W（瓦）、KW（千瓦）、MW（兆瓦）、GW，如果是KW时将数值乘以e3，MW乘以e6，GW乘以e9，提取值为v2；
3. 天线增益G：如果数值没有带dB，需将其转为dB，公式为：dB=10*np.log10(G)，提取值为v3；
4. 雷达截面积sigma：名称还可能为RCS，单位为平方米，$m^{2}$，提取值为v4；
5. 有效噪声温度T_e：名称还可能为噪声温度或工作温度，单位为K（开尔文），如果是摄氏度或度，需将数值加上275.12，提取值为v5；
6. 噪声系数F：单位为dB，若不带dB，将其转为dB，公式为：dB=10*np.log10(F)，提取值为v6；
7. 系统损耗L：单位为dB，若不带dB，将其转为dB，公式为：dB=10*np.log10(L)，提取值为v7；
8. 距离范围Rs：为一组数字，单位为m（米）、Km（千米），如果是Km则将数值乘以e3,，提取值为[v81, v82, ...]；
9. 信噪比范围SNRs：由下限、上限和等分数组成，单位为dB，如果不带dB，需转为dB，转换公式为：dB=10*np.log10(SNR)，提取值为[v91, v92, v93]；
10. 生成json结果：
{
    "f": v1,
    "p_t": v2,
    "G": v3,
    "sigma": v4,
    "T_e": v5,
    "F": v6,
    "L": v7,
    "Rs": [v81, v82, ...],
    "SNRs": [v91, v92, v93]
}

# 示例
输入：假设一个工作频率为$5.6\\times 10^{9}\\text{Hz}$，发射功率为$1.0\\times 10^{6}\\text{W} $，天线增益为40.0dB，雷达截面积为$0.1\\text{m}^{2}$，有效噪声温度为300.0K，噪声系数为5.0dB，系统损耗为6.0dB，距离范围为：75Km、100Km、150Km，研究信噪比范围为5dB到20dB，将其分为200等份，求出对应的最小脉冲宽度并绘制曲线。
输出：
{
    "f": 5.6e9,
    "p_t": 1.0e6,
    "G": 40.0,
    "sigma": 0.1,
    "T_e": 300.0,
    "F": 5.0,
    "L": 6.0,
    "Rs": [75e3, 100e3, 150e3],
    "SNRs": [5, 20, 200]
}

# 当前输入
'''
        prompt = f'{template}{user_query}'
        return prompt